import numpy as np
from tqdm import trange

from .tensor import Tensor


def train(model, X_train, y_train, optim, epochs=10, BS=128, lossfn=lambda x: x.mean()):
    Tensor.training = True
    losses, accuracies = [], []
    for _ in range(len(y_train)//BS+1):
        with trange(epochs) as t:
            for _ in t:
                samp = np.random.randint(0, X_train.shape[0], size=BS)

                x = Tensor(X_train[samp])
                y = y_train[samp]

                # network
                out = model.forward(x)

                loss = lossfn(out, y)
                optim.zero_grad()
                loss.backward()
                optim.step()

                cat = np.argmax(out.data, axis=-1)
                accuracy = (cat == y).mean()

                # printing
                loss = loss.data
                losses.append(loss)
                accuracies.append(accuracy)
                t.set_description("loss %.5f accuracy %.5f" %
                                  (loss, accuracy*100))
    return losses, accuracies


def evaluate(model, X_test, y_test, BS=256):
    Tensor.training = False
    sm = 0
    with trange((len(y_test)-1)//BS+1) as t:
        for i in t:
            if i != 0:
                t.set_description('Acc: {:.3f}%'.format(sm/(i*BS)*100))
            out = model.forward(Tensor(X_test[i*BS:(i+1)*BS])).data.argmax(-1)
            sm += (out == y_test[i*BS:(i+1)*BS]).sum()
    return sm/len(y_test)
